"""
可视化工具模块

该模块提供跟踪结果的可视化功能。
主要功能：
1. 绘制跟踪边界框
2. 绘制跟踪轨迹
3. 显示跟踪信息
4. 保存可视化结果

作者：目标跟踪系统
日期：2024
"""

import cv2
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Tuple, Optional, Dict
import time
from PIL import Image, ImageDraw, ImageFont
import platform


class TrackingVisualizer:
    """
    跟踪可视化器类
    
    提供跟踪结果的可视化功能，包括边界框绘制、轨迹显示等。
    
    参数：
        bbox_color: 边界框颜色，默认为绿色
        trajectory_color: 轨迹颜色，默认为红色
        text_color: 文本颜色，默认为白色
        thickness: 线条粗细，默认为2
    """
    
    def __init__(self, bbox_color: Tuple[int, int, int] = (0, 255, 0), 
                 trajectory_color: Tuple[int, int, int] = (0, 0, 255),
                 text_color: Tuple[int, int, int] = (255, 255, 255),
                 thickness: int = 2):
        self.bbox_color = bbox_color
        self.trajectory_color = trajectory_color
        self.text_color = text_color
        self.thickness = thickness
        
        # 轨迹历史
        self.trajectory_history = []
        self.max_trajectory_length = 50
        
        # 性能统计
        self.fps_history = []
        self.frame_times = []
        
        # 中文字体初始化
        self.font = self._get_chinese_font()
        
    def _get_chinese_font(self):
        """
        获取中文字体
        
        返回：
            PIL字体对象
        """
        try:
            # Windows系统字体路径
            if platform.system() == "Windows":
                font_paths = [
                    "C:/Windows/Fonts/simsun.ttc",  # 宋体
                    "C:/Windows/Fonts/simhei.ttf",  # 黑体
                    "C:/Windows/Fonts/msyh.ttc",    # 微软雅黑
                ]
            # Linux系统字体路径
            elif platform.system() == "Linux":
                font_paths = [
                    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
                    "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
                ]
            # macOS系统字体路径
            else:
                font_paths = [
                    "/System/Library/Fonts/PingFang.ttc",
                    "/System/Library/Fonts/STHeiti Light.ttc",
                ]
            
            # 尝试加载字体，增大字体大小
            for font_path in font_paths:
                try:
                    return ImageFont.truetype(font_path, 40)  # 从20增加到40，放大一倍
                except:
                    continue
                    
            # 如果都失败，使用默认字体
            return ImageFont.load_default()
            
        except Exception:
            return ImageFont.load_default()
            
    def _draw_chinese_text(self, image: np.ndarray, text: str, position: Tuple[int, int], 
                          color: Tuple[int, int, int] = (255, 255, 255)) -> np.ndarray:
        """
        在图像上绘制中文文本
        
        参数：
            image: 输入图像
            text: 要绘制的文本
            position: 文本位置 (x, y)
            color: 文本颜色
            
        返回：
            绘制了文本的图像
        """
        # 将OpenCV图像转换为PIL图像
        pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_image)
        
        # 绘制文本 - PIL的颜色格式是RGB，需要转换
        rgb_color = (color[2], color[1], color[0])  # BGR转RGB
        draw.text(position, text, font=self.font, fill=rgb_color)
        
        # 转换回OpenCV格式
        return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        
    def draw_bbox(self, image: np.ndarray, bbox: Tuple[int, int, int, int], 
                  label: str = "", confidence: float = 0.0, 
                  color: Optional[Tuple[int, int, int]] = None) -> np.ndarray:
        """
        在图像上绘制边界框
        
        参数：
            image: 输入图像
            bbox: 边界框 (x, y, width, height)
            label: 标签文本
            confidence: 置信度
            color: 边界框颜色，如果为None则使用默认颜色
            
        返回：
            绘制了边界框的图像
        """
        if color is None:
            color = self.bbox_color
            
        x, y, w, h = bbox
        
        # 绘制边界框
        cv2.rectangle(image, (x, y), (x + w, y + h), color, self.thickness)
        
        # 绘制标签
        if label or confidence > 0:
            text = f"{label}" if label else ""
            if confidence > 0:
                text += f" ({confidence:.2f})"
                
            if text:
                # 计算文本尺寸
                font = cv2.FONT_HERSHEY_SIMPLEX
                font_scale = 0.6
                text_thickness = 1
                
                (text_width, text_height), baseline = cv2.getTextSize(text, font, font_scale, text_thickness)
                
                # 绘制文本背景
                cv2.rectangle(image, (x, y - text_height - 10), 
                             (x + text_width, y), color, -1)
                
                # 绘制文本
                cv2.putText(image, text, (x, y - 5), font, font_scale, 
                           self.text_color, text_thickness)
                           
        return image
        
    def draw_trajectory(self, image: np.ndarray, trajectory: List[Tuple[float, float]], 
                       color: Optional[Tuple[int, int, int]] = None) -> np.ndarray:
        """
        在图像上绘制跟踪轨迹
        
        参数：
            image: 输入图像
            trajectory: 轨迹点列表 [(x, y), ...]
            color: 轨迹颜色，如果为None则使用默认颜色
            
        返回：
            绘制了轨迹的图像
        """
        if color is None:
            color = self.trajectory_color
            
        if len(trajectory) < 2:
            return image
            
        # 绘制轨迹线
        points = np.array(trajectory, dtype=np.int32)
        cv2.polylines(image, [points], False, color, self.thickness)
        
        # 绘制轨迹点
        for i, point in enumerate(trajectory):
            x, y = int(point[0]), int(point[1])
            
            # 根据时间远近调整点的大小和透明度
            alpha = max(0.3, 1.0 - i / len(trajectory))
            radius = max(2, int(5 * alpha))
            
            # 绘制轨迹点
            cv2.circle(image, (x, y), radius, color, -1)
            
        return image
        
    def update_trajectory(self, center: Tuple[float, float]) -> None:
        """
        更新轨迹历史
        
        参数：
            center: 新的中心点位置
        """
        self.trajectory_history.append(center)
        
        # 限制轨迹长度
        if len(self.trajectory_history) > self.max_trajectory_length:
            self.trajectory_history.pop(0)
            
    def draw_tracking_info(self, image: np.ndarray, tracking_info: Dict) -> np.ndarray:
        """
        在图像上绘制跟踪信息
        
        参数：
            image: 输入图像
            tracking_info: 跟踪信息字典
            
        返回：
            绘制了信息的图像
        """
        line_height = 30  # 调整行高
        padding = 15  # 内边距
        
        # 准备要显示的信息
        info_lines = []
        
        if 'frame_count' in tracking_info:
            info_lines.append(f"帧数: {tracking_info['frame_count']}")
            
        if 'success_rate' in tracking_info:
            info_lines.append(f"成功率: {tracking_info['success_rate']:.2%}")
            
        if 'is_tracking' in tracking_info:
            status = "跟踪中" if tracking_info['is_tracking'] else "跟踪失败"
            info_lines.append(f"状态: {status}")
            
        if 'fps' in tracking_info:
            info_lines.append(f"FPS: {tracking_info['fps']:.1f}")
            
        # 绘制信息背景和文本
        if info_lines:
            # 使用PIL测量实际文本尺寸
            # 将OpenCV图像转换为PIL图像用于测量
            pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(pil_image)
            
            # 计算所有文本的实际宽度和高度
            text_widths = []
            text_heights = []
            
            for line in info_lines:
                bbox = draw.textbbox((0, 0), line, font=self.font)
                text_widths.append(bbox[2] - bbox[0])
                text_heights.append(bbox[3] - bbox[1])
            
            # 计算背景框尺寸
            max_width = max(text_widths) if text_widths else 100
            total_height = sum(text_heights) + (len(info_lines) - 1) * 5  # 行间距5像素
            
            # 背景框坐标
            bg_x1 = 15
            bg_y1 = 15
            bg_x2 = bg_x1 + max_width + 2 * padding
            bg_y2 = bg_y1 + total_height + 2 * padding
            
            # 绘制半透明背景
            overlay = image.copy()
            cv2.rectangle(overlay, (bg_x1, bg_y1), (bg_x2, bg_y2), (0, 0, 0), -1)
            cv2.addWeighted(overlay, 0.8, image, 0.2, 0, image)
            
            # 绘制文本，确保与背景对齐
            current_y = bg_y1 + padding
            for i, line in enumerate(info_lines):
                text_x = bg_x1 + padding
                text_y = current_y
                
                image = self._draw_chinese_text(image, line, (text_x, text_y), self.text_color)
                current_y += text_heights[i] + 5  # 添加行间距
                           
        return image
        
    def draw_search_window(self, image: np.ndarray, search_window: Tuple[int, int, int, int]) -> np.ndarray:
        """
        绘制搜索窗口
        
        参数：
            image: 输入图像
            search_window: 搜索窗口 (x, y, width, height)
            
        返回：
            绘制了搜索窗口的图像
        """
        x, y, w, h = search_window
        
        # 使用虚线绘制搜索窗口
        color = (255, 255, 0)  # 黄色
        thickness = 1
        
        # 绘制虚线矩形
        self._draw_dashed_rectangle(image, (x, y), (x + w, y + h), color, thickness)
        
        return image
        
    def _draw_dashed_rectangle(self, image: np.ndarray, pt1: Tuple[int, int], 
                              pt2: Tuple[int, int], color: Tuple[int, int, int], 
                              thickness: int, dash_length: int = 10) -> None:
        """
        绘制虚线矩形
        
        参数：
            image: 输入图像
            pt1: 左上角点
            pt2: 右下角点
            color: 颜色
            thickness: 线条粗细
            dash_length: 虚线长度
        """
        x1, y1 = pt1
        x2, y2 = pt2
        
        # 绘制四条虚线边
        self._draw_dashed_line(image, (x1, y1), (x2, y1), color, thickness, dash_length)  # 上边
        self._draw_dashed_line(image, (x2, y1), (x2, y2), color, thickness, dash_length)  # 右边
        self._draw_dashed_line(image, (x2, y2), (x1, y2), color, thickness, dash_length)  # 下边
        self._draw_dashed_line(image, (x1, y2), (x1, y1), color, thickness, dash_length)  # 左边
        
    def _draw_dashed_line(self, image: np.ndarray, pt1: Tuple[int, int], 
                         pt2: Tuple[int, int], color: Tuple[int, int, int], 
                         thickness: int, dash_length: int) -> None:
        """
        绘制虚线
        
        参数：
            image: 输入图像
            pt1: 起点
            pt2: 终点
            color: 颜色
            thickness: 线条粗细
            dash_length: 虚线长度
        """
        x1, y1 = pt1
        x2, y2 = pt2
        
        # 计算线段长度和方向
        dx = x2 - x1
        dy = y2 - y1
        length = np.sqrt(dx * dx + dy * dy)
        
        if length == 0:
            return
            
        # 单位方向向量
        unit_x = dx / length
        unit_y = dy / length
        
        # 绘制虚线段
        current_pos = 0
        while current_pos < length:
            # 计算当前段的起点和终点
            start_x = int(x1 + current_pos * unit_x)
            start_y = int(y1 + current_pos * unit_y)
            
            end_pos = min(current_pos + dash_length, length)
            end_x = int(x1 + end_pos * unit_x)
            end_y = int(y1 + end_pos * unit_y)
            
            # 绘制实线段
            cv2.line(image, (start_x, start_y), (end_x, end_y), color, thickness)
            
            # 跳过间隔
            current_pos += dash_length * 2
            
    def calculate_fps(self, frame_time: float) -> float:
        """
        计算帧率
        
        参数：
            frame_time: 当前帧处理时间
            
        返回：
            当前帧率
        """
        self.frame_times.append(frame_time)
        
        # 保持最近30帧的时间记录
        if len(self.frame_times) > 30:
            self.frame_times.pop(0)
            
        # 计算平均帧率
        if len(self.frame_times) > 1:
            avg_time = np.mean(self.frame_times)
            fps = 1.0 / avg_time if avg_time > 0 else 0
        else:
            fps = 0
            
        return fps
        
    def create_result_image(self, image: np.ndarray, bbox: Optional[Tuple[int, int, int, int]], 
                           tracking_info: Dict, show_trajectory: bool = True) -> np.ndarray:
        """
        创建包含所有可视化信息的结果图像
        
        参数：
            image: 输入图像
            bbox: 目标边界框
            tracking_info: 跟踪信息
            show_trajectory: 是否显示轨迹
            
        返回：
            完整的可视化结果图像
        """
        result_image = image.copy()
        
        # 绘制边界框
        if bbox is not None:
            confidence = tracking_info.get('similarity', 0.0)
            color = self.bbox_color if tracking_info.get('is_tracking', False) else (0, 0, 255)
            result_image = self.draw_bbox(result_image, bbox, "目标", confidence, color)
            
            # 更新轨迹
            center_x = bbox[0] + bbox[2] // 2
            center_y = bbox[1] + bbox[3] // 2
            self.update_trajectory((center_x, center_y))
            
        # 绘制轨迹
        if show_trajectory and len(self.trajectory_history) > 1:
            result_image = self.draw_trajectory(result_image, self.trajectory_history)
            
        # 绘制跟踪信息
        result_image = self.draw_tracking_info(result_image, tracking_info)
        
        return result_image
        
    def save_result(self, image: np.ndarray, filename: str) -> bool:
        """
        保存可视化结果
        
        参数：
            image: 要保存的图像
            filename: 文件名
            
        返回：
            是否成功保存
        """
        try:
            cv2.imwrite(filename, image)
            return True
        except Exception as e:
            print(f"保存图像失败: {e}")
            return False
            
    def reset_trajectory(self) -> None:
        """
        重置轨迹历史
        """
        self.trajectory_history.clear()
        
    def set_colors(self, bbox_color: Optional[Tuple[int, int, int]] = None,
                   trajectory_color: Optional[Tuple[int, int, int]] = None,
                   text_color: Optional[Tuple[int, int, int]] = None) -> None:
        """
        设置可视化颜色
        
        参数：
            bbox_color: 边界框颜色
            trajectory_color: 轨迹颜色
            text_color: 文本颜色
        """
        if bbox_color is not None:
            self.bbox_color = bbox_color
        if trajectory_color is not None:
            self.trajectory_color = trajectory_color
        if text_color is not None:
            self.text_color = text_color 